﻿#ifdef _DEBUG
#   include <WinSock2.h>
#   include <vld.h>
#endif

#include <KDepends.h>
#include <KUtility.h>
#include <KWinSock.h>
#include <iomanip>

#ifdef _DEBUG
#   pragma comment(lib, "Kirnod.lib")
#else
#   pragma comment(lib, "Kirno.lib")
#endif

#include "S5ActionClient.h"
#include "S5Def.h"
#include "S5Server.h"
#include "S5Utility.h"

using namespace std;

K_ANONYMOUS_NAMESPACE_BEG

const char *    APP_NAME = "Kirno Socks5 Server";

const char *    CMD_STR[] = {
    "UNKNOWN", "CONNECT", "BIND", "UDP ASSOCIATE"
};

const char *    STATE_STR[] = {
    "TIMEOUT",
    "VER1", "NMETHODS", "METHODS",
    "VER2", "ULEN", "UNAME", "PLEN", "PASSWD",
    "VER3", "CMD", "RSV", "ATYP", "DLEN", "DST.ADDR", "DST.PORT",
    "EVAL",
    "TCP RELAY", "UDP RELAY"
};

const char *    UNIT_STR[] = {
    "KB", "MB", "GB", "TB"
};

KLock           g_output;
KManualEvent    g_shutdown;
KNetScene *     g_scene_dbg = NULL;
KNetScene *     g_scene_svr = NULL;
KTaskProgress * g_progress = NULL;
OutputProxy     g_debug;

inline void adjustUnit(double &kbps, size_t &unit)
{
    do {
        if (kbps >= 1000) {
            kbps /= 1024;
            ++unit;
        } else {
            return;
        }
    } while (unit + 1 < sizeof(UNIT_STR) / sizeof(UNIT_STR[0]));
}

class ServerWatcher : public KTimer {
public:
    enum Mode { Default = 0, Performance, Health, Pool, Speed };

public:
    ServerWatcher() : m_mode(ServerWatcher::Default) {}

public:
    int mode() const { return m_mode; }

    void switchMode()
    {
        if (m_mode == ServerWatcher::Speed) {
            m_mode = ServerWatcher::Default;
        } else {
            ++m_mode;
        }
        tick();
    }

protected:
    void tick();

private:
    int m_mode;
    KUtility::cpu_rate_cache_t m_cache;
} g_watcher;

class Socks5Server : public S5Server {
public:
    explicit Socks5Server(LONG max = 10000, u_short port = 1080)
        : S5Server(max, port), m_total(extraAccepts() + 1) {}

public:
    LONG totalAccepts() const { return m_total; }

public:
    static BOOL WINAPI handler(DWORD dwEvent)
    {
        switch (dwEvent) {
        case CTRL_BREAK_EVENT:
            g_watcher.switchMode();
            break;
        case CTRL_C_EVENT: case CTRL_LOGOFF_EVENT: case CTRL_SHUTDOWN_EVENT: case CTRL_CLOSE_EVENT:
            g_scene_svr->stopAll();
            g_shutdown.set();
            break;
        default:
            return FALSE;
        }
        return TRUE;
    }

protected:
    // 如果需要限制连接的客户端 IP, 请实现 allow
    // 如果需要用户名/密码的身份验证, 请实现 requireAuthentication 和 verify

    // 说明: 控制台窗口的输出并不是必须的, 若需要提高服务器运行的效率和稳定性, 请自行关闭(全部都在本文件中)
    // 注意: 按住控制台窗口的滚动条或进入快速编辑模式会导致正在输出的线程被阻塞, 此后程序的行为无定义(多为死锁)

    void onAborted(S5ActionClient *client)
    {
        KMutex lock(g_output);
        if (client->state() > S5ActionClient::EVAL) {
            if (client->username().length()) {
                g_debug << OutputProxy::DarkGreen << "Client "
                    << OutputProxy::Gray << client->username()
                    << OutputProxy::Dark << '(' << client->address().ip() << ':' << client->address().port() << ')';
            } else {
                g_debug << OutputProxy::DarkGreen << "Client "
                    << OutputProxy::Gray << client->address().ip() << ':' << client->address().port();
            }
            g_debug << OutputProxy::DarkGreen << " finished "
                << OutputProxy::DarkYellow << CMD_STR[client->command()] << "\r\n" << OutputProxy::Flush;
        } else if (client->state() == S5ActionClient::EVAL) {
            if (client->username().length()) {
                g_debug << OutputProxy::DarkRed << "Client "
                    << OutputProxy::Gray << client->username()
                    << OutputProxy::Dark << '(' << client->address().ip() << ':' << client->address().port() << ')';
            } else {
                g_debug << OutputProxy::DarkRed << "Client "
                    << OutputProxy::Gray << client->address().ip() << ':' << client->address().port();
            }
            int cmd = client->command();
            if (cmd == S5_CMD_UDP_ASSOCIATE) {
                g_debug << OutputProxy::DarkRed << " failed "
                    << OutputProxy::DarkYellow << CMD_STR[cmd] << "\r\n" << OutputProxy::Flush;
            } else {
                g_debug << OutputProxy::DarkRed << " failed "
                    << OutputProxy::DarkYellow << CMD_STR[cmd]
                << OutputProxy::Dark << ' ' << client->domain() << "\r\n" << OutputProxy::Flush;
            }
        } else {
            g_debug << OutputProxy::DarkRed << "Client "
                << OutputProxy::Gray << client->address().ip() << ':' << client->address().port()
                << OutputProxy::DarkRed << " has aborted in "
                << OutputProxy::DarkCyan << STATE_STR[client->state()]
                << OutputProxy::DarkRed << " state" << "\r\n" << OutputProxy::Flush;
        }
    }

    void onClientsChanged(LONG current)
    {
        switch (current) {
        case 0:
            if (extraAccepts() >= 0) {
                Console::setIcon(Console::Idle);
            } else {
                Console::setIcon(Console::Dead);
            } break;
        case 1:
            Console::setIcon(Console::Working);
            break;
        }
    }

    void onEval(S5ActionClient *client)
    {
        KMutex lock(g_output);
        if (client->username().length()) {
            g_debug << OutputProxy::Green << "Client "
                << OutputProxy::White << client->username()
                << OutputProxy::Dark << '(' << client->address().ip() << ':' << client->address().port() << ')';
        } else {
            g_debug << OutputProxy::Green << "Client "
                << OutputProxy::White << client->address().ip() << ':' << client->address().port();
        }
        int cmd = client->command();
        if (cmd == S5_CMD_UDP_ASSOCIATE) {
            g_debug << OutputProxy::Green << " started "
                << OutputProxy::Yellow << CMD_STR[cmd] << "\r\n" << OutputProxy::Flush;
        } else {
            g_debug << OutputProxy::Green << " started "
                << OutputProxy::Yellow << CMD_STR[cmd]
            << OutputProxy::Dark << ' ' << client->domain() << "\r\n" << OutputProxy::Flush;
        }
    }

    void onStopping()
    {
        // Test only
        g_debug << OutputProxy::Gray << "Server exiting..." << "\r\n" << OutputProxy::Flush;
    }

private:
    const LONG m_total;
} g_server(512); // 服务器最多允许的连接数, 使用默认端口 1080

void ServerWatcher::tick()
{
    switch (m_mode) {
    case ServerWatcher::Default:
        {
            if (g_debug.portUsed()) {
                ostringstream oss;
                oss << APP_NAME << " (Port: " << g_debug.portUsed() << ')';
                g_debug.setTitle(oss.str());
            } else {
                g_debug.setTitle(APP_NAME);
            }
            g_progress->setStatus(g_server.clients() > 0 ? TBPF_INDETERMINATE : TBPF_NOPROGRESS);
        } break;
    case ServerWatcher::Performance:
        {
            int cpu = 100;
            uint64_t mem;
            ostringstream oss;
            if (KUtility::getProcessCpuRate(m_cache, &cpu)) {
                oss << cpu << "%, ";
            } else {
                oss << "N/A, ";
            }
            if (KUtility::getProcessMemUsage(&mem)) {
                oss << (mem >> 10) << "KB - ";
            } else {
                oss << "N/A - ";
            }
            oss << APP_NAME;
            g_debug.setTitle(oss.str());
            if (cpu > 0) {
                g_progress->setMaxProgress(100);
                g_progress->setStatus(TBPF_ERROR);
                g_progress->setProgress(cpu);
            } else {
                g_progress->setStatus(TBPF_NOPROGRESS);
            }
        } break;
    case ServerWatcher::Health:
        {
            LONG current = g_server.clients();
            if (current > g_server.maxClients()) {
                current = g_server.maxClients();
            }
            ostringstream oss;
            oss << "Clients: " << current << '/' << g_server.maxClients()
                << ", Total: " << g_server.totalClients()
                << ", Activity: " << g_server.extraAccepts() + 1 << '/' << g_server.totalAccepts()
                << " - " << APP_NAME;
            g_debug.setTitle(oss.str());
            if (current > 0) {
                g_progress->setMaxProgress(g_server.maxClients());
                g_progress->setStatus(TBPF_PAUSED);
                g_progress->setProgress(current);
            } else {
                g_progress->setStatus(TBPF_NOPROGRESS);
            }
        } break;
    case ServerWatcher::Pool:
        {
            ostringstream oss;
            oss << "Item: " << g_server.poolSize() << '/' << g_server.maxPoolSize()
                << ", Scene: " << g_scene_svr->poolSize() << '/' << g_scene_svr->maxPoolSize()
                << " - " << APP_NAME;
            g_debug.setTitle(oss.str());
            if (g_scene_svr->poolSize() > 0) {
                g_progress->setMaxProgress(g_scene_svr->maxPoolSize());
                g_progress->setStatus(TBPF_NORMAL);
                g_progress->setProgress(g_scene_svr->poolSize());
            } else {
                g_progress->setStatus(TBPF_NOPROGRESS);
            }
        } break;
    default:
        {
            double bytes_down(g_server.kbDownload()), bytes_up(g_server.kbUpload());
            double speed_down(g_server.kbpsDownload()), speed_up(g_server.kbpsUpload());
            size_t unit_bytes_down = 0, unit_bytes_up = 0;
            size_t unit_speed_down = 0, unit_speed_up = 0;
            ULONGLONG current_speed_down = ULONGLONG(speed_down);
            static ULONGLONG max_speed_down = 0;
            adjustUnit(bytes_down, unit_bytes_down);
            adjustUnit(bytes_up, unit_bytes_up);
            adjustUnit(speed_down, unit_speed_down);
            adjustUnit(speed_up, unit_speed_up);
            if (max_speed_down < current_speed_down) {
                max_speed_down = current_speed_down;
            }
            ostringstream oss;
            oss << setiosflags(ios_base::fixed) << setprecision(2)
                << "Down: " << bytes_down << UNIT_STR[unit_bytes_down] << " at "
                << speed_down << UNIT_STR[unit_speed_down]
                << "/s, Up: " << bytes_up << UNIT_STR[unit_bytes_up] << " at "
                << speed_up << UNIT_STR[unit_speed_up]
                << "/s - " << APP_NAME;
                g_debug.setTitle(oss.str());
            if (current_speed_down > 0) {
                g_progress->setMaxProgress(max_speed_down);
                g_progress->setStatus(TBPF_NORMAL);
                g_progress->setProgress(current_speed_down);
            } else {
                g_progress->setStatus(TBPF_NOPROGRESS);
            }
        } break;
    }
}

K_ANONYMOUS_NAMESPACE_END

int main(int argc, char **argv)
{
    KWinSock winsock;
    CoInitialize(NULL);
    Console::loadIcon(argv[0]);
#ifndef _DEBUG
    AllocConsole();
#endif
    Console::setIcon(Console::Idle);
    {
        HWND hConsole = GetConsoleWindow();
        if (IsWindow(hConsole)) {
            HMENU hMenu = GetSystemMenu(hConsole, FALSE);
            if (IsMenu(hMenu)) {
                RemoveMenu(hMenu, SC_CLOSE, MF_BYCOMMAND);
            }
        }
        g_progress = new KTaskProgress(hConsole);
        g_scene_dbg = new KNetScene(KUtility::numOfProcessors() * 2);
        g_scene_svr = new KNetScene(KUtility::numOfProcessors() * 2, g_server.maxClients() * 4);
        g_watcher.start(0, 1000);
        SetConsoleCtrlHandler(Socks5Server::handler, TRUE);
        {
            g_server.setMaxPoolSize(g_server.maxClients() * 2);
            g_scene_dbg->install(&g_debug);
            g_scene_dbg->start();
            g_scene_svr->install(&g_server);
            g_scene_svr->start();
            if (g_server.isRunning()) {
                g_shutdown.wait();
            } else {
                g_scene_svr->stopAll();
                g_debug << OutputProxy::Gray << "Server failed to start" << "\r\n" << OutputProxy::Flush;
            }
        }
        SetConsoleCtrlHandler(Socks5Server::handler, FALSE);
        g_watcher.stop();
        delete g_scene_svr;
        delete g_scene_dbg;
        delete g_progress;
    }
#ifndef _DEBUG
    FreeConsole();
#endif
    Console::freeIcon();
    CoUninitialize();
    return 0;
}